#include "kernel_operator.h"
#include "math.h"

constexpr int32_t TOTAL_LENGTH = 2048;
constexpr int32_t USE_CORE_NUM = 8;
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM;
constexpr int32_t TILE_NUM = 8;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM;

using namespace AscendC;

template <typename T>
class CBRTKernalFast
{

private:
    TQue<QuePosition::VECIN, BUFFER_NUM> inX1;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outY;
    TBuf<QuePosition::VECCALC> tmp2buf, tmp3buf;
    GlobalTensor<T> x1Gm, yGm;

    uint32_t L, R;
    TPipe pipe;

#define BUF_SZ 64

public:
    __aicore__ inline CBRTKernalFast() {}

    __aicore__ inline void Init(GM_ADDR input, GM_ADDR out, uint32_t size, uint32_t length)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        const unsigned num_cores = GetBlockNum();
        //this->pipe = PIPE;

        // printf("size:%d length:%d\n",size,length);

        unsigned L = GetBlockIdx() * length;
        unsigned R = L + length;
        if (R > size)
            R = size;

        x1Gm.SetGlobalBuffer((__gm__ T *)input + L,length);
        yGm.SetGlobalBuffer((__gm__ T *)out + L,length);
        this->R = R - L;
        this->L = 0;

        pipe.InitBuffer(inX1, BUFFER_NUM, BUF_SZ * sizeof(T));
        pipe.InitBuffer(outY, BUFFER_NUM, BUF_SZ * sizeof(T));
        pipe.InitBuffer(tmp2buf, BUF_SZ * sizeof(T));
        pipe.InitBuffer(tmp3buf, BUF_SZ * sizeof(T));
    }
    __aicore__ inline void CopyIn(int i, int sz)
    {
        LocalTensor<T> src0Local = inX1.AllocTensor<T>();
        DataCopy(src0Local, x1Gm[i], sz);
        inX1.EnQue(src0Local);
    }
    __aicore__ inline void Compute(int i, int sz)
    {
        LocalTensor<T> x1Local = inX1.DeQue<T>();
        LocalTensor<T> yLocal = outY.AllocTensor<T>();
        LocalTensor<T> tmp2 = tmp2buf.Get<T>();
        LocalTensor<T> tmp3 = tmp3buf.Get<T>();

        DataCopy(yLocal, x1Local, sz);

        /*
               z = z - \frac{z^3 - A}{3z^2}*/
        for (int iter = 0; iter < 128; iter++)
        {
            Mul(tmp2, yLocal, yLocal, sz); // tmp2 = y^2
            Mul(tmp3, yLocal, tmp2, sz);   // tmp3 = y^3
            Muls(tmp2, tmp2, (T)3, sz);    // tmp2 = 3 * y^2
            Sub(tmp3, tmp3, x1Local, sz);  // tmp3 = y^3 - x1
            Div(tmp3, tmp3, tmp2, sz);     // tmp3 = (y^3 - x1) / (3 * y^2)
            Sub(yLocal, yLocal, tmp3, sz); // y = y - (y^3 - x1) / (3 * y^2)
        }


        tmp2buf.FreeTensor(tmp2);
        tmp3buf.FreeTensor(tmp3);
        inX1.FreeTensor(x1Local);
        outY.EnQue<T>(yLocal);
    }

    __aicore__ inline void CopyOut(int i, int sz)
    {
        LocalTensor<T> yoLocal = outY.DeQue<T>();
        DataCopy(yGm[i], yoLocal, sz);
        outY.FreeTensor(yoLocal);
    }

    __aicore__ inline void Process()
    {
        uint32_t i = 0;
        for (; i < R; i += BUF_SZ)
        {
            CopyIn(i, BUF_SZ);
            Compute(i, BUF_SZ);
            CopyOut(i, BUF_SZ);
        }
        /*if (i < R)
        {
            uint32_t sz = R - i;
            CopyIn(i, sz);
            Compute(i, BUF_SZ);
            CopyOut(i, sz);
        }*/
    }
};

extern "C" __global__ __aicore__ void cbrt_custom(GM_ADDR input, GM_ADDR out)
{
    // GET_TILING_DATA(tiling_data, tiling);
    
    CBRTKernalFast<float> op;
    op.Init(input, out, TOTAL_LENGTH, BLOCK_LENGTH);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void cbrt_custom_do(uint32_t blockDim, void *stream, float *x, float *z)
{
    cbrt_custom<<<blockDim, nullptr, stream>>>(x, z);
}
#endif
